from typing import Dict, Union
import tornado
import tornado.websocket


class BaseWebSocketContext:
    """
    TornadoWebSocketContext is a class that stores the context of a Tornado WebSocket connection.
    """

    def send_message(self, message: Union[str, Dict]):
        raise NotImplementedError

    def close(self):
        raise NotImplementedError


class TornadoWebSocketContext(BaseWebSocketContext):

    def __init__(
        self,
        event_loop: tornado.ioloop.IOLoop,
        ws: tornado.websocket.WebSocketHandler,
    ):
        super().__init__()
        self.event_loop = event_loop
        self.ws = ws

    def send_message(self, message: Union[str, Dict]):
        self.event_loop.add_callback(self.ws.write_message, message)

    def close(self):
        self.ws.close()


class BaseWebSocketHandler:
    def __init__(self, ctx: BaseWebSocketContext) -> None:
        self.ctx = ctx

    def on_open(self, params: Dict[str, str] = None):
        pass

    def on_message(self, message):
        pass

    def on_close(self):
        pass
